import numpy as np
import os
class LinearIntervals(object):
    def __init__(self):
        self.a0 = None 
        self.a1 = None
        self.x  = None
        self.y  = None

    def calc_coeffs( self, x, y ):
        self.x = x
        self.y = y
        self.a0 = np.zeros(len(x)-1,dtype='float16')
        self.a1 = np.zeros(len(x)-1,dtype='float16')
        for i in range( len(x)-1 ):
            y1 = y[i]
            y2 = y[i+1]
            x1 = x[i]
            x2 = x[i+1]
            self.a1[i] = (y2-y1)/(x2-x1)
            self.a0[i] = (x2*y1-x1*y2)/(x2-x1) 
    def aux_calc_x( self, x):
        if x < self.x[0]:
            return self.a0[0] + self.a1[0]*x
        for i in range( len(self.x) - 1 ): 
            xi  = self.x[i]
            xi1 = self.x[i+1]
            if x >= xi and x < xi1:
                return self.a0[i] + self.a1[i]*x
                break
        if x >= self.x[-1]:
            return self.a0[-1] + self.a1[-1]*x
        
    def calc_y( self, x ):
        if x.__class__.__name__ == 'list'\
        or x.__class__.__name__.find('array') > -1:
            y = np.zeros(len(x),dtype='float16')
            for i in range(len(x)):
                xi   = x[i]
                y[i] = self.aux_calc_x( xi )
            return y        
        else:
            return self.aux_calc_x( x )

class Function(object):
    def __init__(self):
        self.name = ''
        self.plot_interval = 40
        self.subplot = None
        self.x = []
        self.y = []
        self.xfit = []
        self.ymins = []
        self.ymaxs = []
        self.poly_min = None
        self.poly_max = None
        self.xplot = None
        self.yminplot = None
        self.ymaxplot = None

    def rebuild(self):
        pass

    def fit_data( self ):
        self.calc_ymins_ymaxs()
        self.poly_min = LinearIntervals()
        self.poly_min.calc_coeffs( self.xfit, self.ymins )
        self.poly_max = LinearIntervals()
        self.poly_max.calc_coeffs( self.xfit, self.ymaxs )
        use_xfit = False
        if not use_xfit:
            p = float(self.plot_interval)
            if self.name[:2] == 'f_':
                self.xplot = np.array([ i*1./p for i in range(p + 1) ])
            else:
                self.xplot = np.array([ i*1./p for i in range(-p, p + 1) ])
        else:
            self.xplot = self.xfit
        self.yminplot = self.poly_min.calc_y( self.xplot ) 
        self.ymaxplot = self.poly_max.calc_y( self.xplot )

    def calc_ymins_ymaxs( self ):
        ymin = {}
        ymax = {}
        for i in range( len(self.x) ):
            xi = self.x[i]
            yi = self.y[i]
            key = float('%2.6f' % xi)
            if key in ymin.keys():
                if yi < ymin[key]:
                    ymin[key] = yi
            else:
                ymin[key] = yi
            if key in ymax.keys():
                if yi > ymax[key]:
                    ymax[key] = yi
            else:
                ymax[key] = yi
        self.xfit = []
        self.ymins = []
        self.ymaxs = []
        sortk = [k for k in ymin.keys()]
        sortk.sort()
        for k in sortk:
            self.xfit.append( k )
            self.ymins.append( ymin[k] )
        sortk = [k for k in ymax.keys()]
        sortk.sort()
        for k in sortk:
            self.ymaxs.append( ymax[k] )

    def calc_limits(self, value):
        return self.poly_min.calc_y( value ),\
               self.poly_max.calc_y( value )


class Parameter(object):
    def __init__(self):
        self.name = ''
        self.keys_default = ['a1','a3','d1','d2','d3','d4',
                                       'b1','b2','b3','b4',
                             'f_00','f_90','f_45','f_ang'   ]
        self.value = None
        self.fig = None
        self.functions = {}
        self.family = None
        self.fig = None
        self.mng = None
        self.constr_order =\
            ['a2','a4','a1','a3','d1','d3','d2','d4','b1','b2','b3','b4']
        self.constr = { 'a1':None, 'a2':None, 'a3':None, 'a4':None,
                        'd1':None, 'd2':None, 'd3':None, 'd4':None,
                        'b1':None, 'b2':None, 'b3':None, 'b4':None  }    

    def calc_feasible_range( self ):
        ans_minv = -1.
        ans_maxv =  1.
        for key in self.constr_order:
            cons_v = self.constr[ key ]
            if cons_v <> None:
                minv,maxv = self.functions[ key ].calc_limits( cons_v )
                if minv > ans_minv:
                    ans_minv = minv
                if maxv < ans_maxv:
                    ans_maxv = maxv
        return ans_minv, ans_maxv    

    def calc_constr_value( self, f ):
        minv, maxv = self.calc_feasible_range()
        ans = (f+1) * (minv-maxv)/2 + minv 
        return ans

    def rebuild( self ):
        for func in self.functions.values():
            func.rebuild()

    def plot( self, keys = 'default',
                    savefig=True,
                    ext='png',
                    fit_plots=False ):
        import matplotlib.pyplot as plt
        if self.fig == None:
            self.fig = plt.figure()
        if self.family == None:
            self.fig.canvas.set_window_title('parameter_%s' % self.name)
        else:
            self.fig.canvas.set_window_title('%s_parameter_%s' \
                                             % (self.family.name, self.name) )
        if keys == 'default':
            keys = self.keys_default
        elif keys == 'all':
            keys = self.functions.keys()
        else:
            pass # user supplied list of keys
        for key in keys:
            if key <> self.name\
            and not key in self.functions.keys():
                keys.pop( keys.index(key) )
        nums = [1,4,9,16,25,36,49,64,81]
        for i in range( 1, len( nums ) ):
            num_1 = nums[ i-1 ]
            num   = nums[  i  ]
            if len( keys ) > num_1 and len( keys ) <= num:
                row = num**0.5
                col = num**0.5
                break
        keys.sort()
        for i in range( len( keys ) ):
            key = keys[i]
            if key == self.name:
                continue
            func = self.functions[ key ] 
            func.subplot = self.fig.add_subplot( row, col, i+1 )
            func.subplot.set_xlabel( func.name )
            #func.subplot.set_xticklabels([])
            if func.name[:2] == 'f_':
                func.subplot.set_xlim(-0.1,1.1)
            else:
                func.subplot.set_xlim(-1.1,1.1)
            if self.name[:2] == 'f_':
                func.subplot.set_ylim(-0.1,1.1)
            else:
                func.subplot.set_ylim(-1.1,1.1)
            if fit_plots:
                func.fit_data()
                func.subplot.scatter( func.xfit, func.ymins,
                                      marker='.' )
                func.subplot.scatter( func.xfit, func.ymaxs,
                                      marker='.' )
                func.subplot.plot( func.xplot, func.yminplot )
                func.subplot.plot( func.xplot, func.ymaxplot )
            else:
                func.subplot.scatter( func.x, func.y,
                                      marker='.' )
        self.fig.show()     
        self.mng = plt.get_current_fig_manager()
        self.mng.window.state('zoomed')
        self.fig.show()     
        if savefig:
            if self.family == None:
                self.fig.savefig(r'c:\Temp\parameter_%s.%s' % (self.name, ext))
            else:
                self.fig.savefig(r'c:\Temp\%s_parameter_%s.%s' \
                                % ( self.family.name, self.name, ext))
            self.mng.window.state('withdrawn')

class LamFamily(object):
    def __init__( self ):
        self.family_name = ''        
        self.num_plies = None
        self.angles = []
        self.lams = []
        self.params = {}
        self.param_names = [ 'a1','a2','a3','a4',
                             'b1','b2','b3','b4',
                             'd1','d2','d3','d4',
                             'f_00','f_90','f_ang',
                             'f_30','f_45','f_60'  ]
    def add_lam( self, lam ):
        lam.family = self
        self.lams.append( lam )

    def rebuild( self ):
        # parameters
        self.params = {}
        for param_name in self.param_names:
            param = Parameter()
            param.family = self
            param.name = param_name
            for func_name in self.param_names:
                if func_name <> param_name:
                    func = Function()
                    func.name = func_name
                    param.functions[ func_name ] = func
            self.params[ param_name ] = param
        # name
        self.name = 'LamFamily_' 
        tmp = ''
        for theta in self.angles:
            self.name += '_%02d_' % theta
        self.name += 'num_plies_%02d' % self.num_plies

    def create_family( self, balanced=True, use_iterator=False ):
        self.rebuild()
        import laminate_stack_calculator
        laminate_stack_calculator = reload( laminate_stack_calculator )
        stacks = laminate_stack_calculator.calc_stacks(
                    angles = self.angles,
                    num_plies = self.num_plies,
                    use_iterator = use_iterator )
        # validating stacks
        self.lams = []
        for stack in stacks:
            lam = Lam()
            lam.family = self
            lam.create_stack( stack, balanced=balanced )
            if lam.valid:
                lam.calc_lam()
                self.add_lam( lam )
        for lam in self.lams:
            lam.calc_lam()
            for param_name, param in self.params.iteritems():
                yi = lam.params[ param_name ]
                for func_name, func in param.functions.iteritems():
                    xi = lam.params[ func_name ]
                    func.x.append( xi )
                    func.y.append( yi )
             
    def plot( self, keys='default',
              savefig=True,
              ext='png',
              fit_plots=False ):
        for param in self.params.values():
            param.plot( keys      = keys       ,
                        savefig   = savefig    , 
                        ext       = ext        ,
                        fit_plots = fit_plots  )
    
    def rebuild_params( self ):
        for param in self.params.values():
            param.rebuild()

    def load( self, path, append=False ):
        import copy
        import cPickle as pickle
        if not os.path.isfile( path ):
            print 'LamFamily pickle %s does not exist!' % path
            return
        tmp = open( path, 'r' )
        tmpfam = pickle.load( tmp )
        tmp.close()
        if not append:
            self = copy.deepcopy( tmpfam )
            del tmpfam
            return self
        else:
            for lam in tmpfam.lams:
                self.lams.append( lam )
            for k,v in tmpfam.params.iteritems():
                self.params[k] = v
            self.name = 'appended_' + self.name 

    def save( self, path ):
        # deleting some objects:
        for param in self.params.values():
            del param.fig
            del param.mng
            for func in param.functions.values():
                del func.subplot

        import cPickle as pickle
        tmp = open( path, 'w' )
        pickle.dump( self, tmp )
        tmp.close()
        print 'LamFamily saved at %s' % path
        return path

class Lam(LamFamily):
    def __init__( self ):
        super( Lam, self ).__init__()
        self.family = None
        self.lam_name = ''
        self.plyt = 0.125
        self.stack = []
        self.valid = False
        self.remaining = None
    
    def check_balance( self ):
        pos={}
        for theta in self.stack:
            theta = int(theta)
            if theta == 0 or theta == 90:
                continue
            if   theta > 0 and not theta in pos.keys():
                pos[ theta ]  = 1
            elif theta > 0 and theta in pos.keys():
                pos[ theta ] += 1
            elif theta < 0 and not abs(theta) in pos.keys():
                pos[ abs(theta) ] = -1
            elif theta < 0 and abs(theta) in pos.keys():
                pos[ abs(theta) ] -= 1
        balanced = True
        for num in pos.values():
            if num <> 0:
                balanced = False
                break
        return balanced

    def check_dist_ang_plies( self, maxdist = 1 ):
        #TODO make this function more general than 30, 45, 60
        p30 = {}
        p45 = {}
        p60 = {}
        for i in range( len(self.stack) ):
            theta = self.stack[i]
            if   theta == 30:
                p30[ i ] = False 
            elif theta == 45:
                p45[ i ] = False
            elif theta == 60:
                p60[ i ] = False
        index_range = range( -maxdist, maxdist+1 )
        index_range.pop( index_range.index(0) )
        for i in range( len(self.stack) ):
            theta = self.stack[i]
            for j in index_range:
                if   theta == -30:
                    if i+j in p30.keys():
                        p30[ i+j ] = True
                elif theta == -45:
                    if i+j in p45.keys():
                        p45[ i+j ] = True
                elif theta == -60:
                    if i+j in p60.keys():
                        p60[ i+j ] = True

        false = False
        if false in p30.values():
            return False
        if false in p45.values():
            return False
        if false in p60.values():
            return False
        return True

    def add_ply( self, theta ):
        if self.remaining > 0:
            if theta == 0 or theta == 90:
                self.remaining -= 1
                self.stack.append(  theta )
            else:
                self.remaining -= 2
                self.stack.append(  theta )
                self.stack.append( -theta )

    def create_stack( self, stack, balanced=True ):
        if not self.family == None:
            self.num_plies = self.family.num_plies
        self.stack = stack
        if balanced:
            self.valid = self.check_balance()
        else:
            self.valid = True
        if self.valid:
            name = '_'.join([ ('%02d' % theta) for theta in stack ])
            self.name  = name
            self.valid = True
            c_ang = 0
            c_00  = 0
            c_30  = 0
            c_45  = 0
            c_60  = 0
            c_90  = 0
            for theta in stack:
                if   theta      == 0:
                    c_00 += 1
                elif theta      == 90:
                    c_90 += 1
                elif abs(theta) == 30:
                    c_30 += 1
                    c_ang += 1
                elif abs(theta) == 45:
                    c_45 += 1
                    c_ang += 1
                elif abs(theta) == 60:
                    c_60 += 1
                    c_ang += 1
                else:
                    c_ang += 1
            self.params[ 'f_00'  ] = float(c_00)  / self.num_plies
            self.params[ 'f_30'  ] = float(c_30)  / self.num_plies
            self.params[ 'f_45'  ] = float(c_45)  / self.num_plies
            self.params[ 'f_60'  ] = float(c_60)  / self.num_plies
            self.params[ 'f_90'  ] = float(c_90)  / self.num_plies
            self.params[ 'f_ang' ] = float(c_ang) / self.num_plies

    def calc_lam( self ):
        import mapy
        import mapy.model.properties.composite as composite
        import mapy.model.materials.matlamina as matlamina
        #laminaProp = ( e1, e2, nu12, g12, g13, g23, tempref )
        laminaprop  = (125.774e3,10.03e3,0.271,5.555e3,5.555e3,3.4e3,273.15)
        m1 = matlamina.read_laminaprop(laminaprop=laminaprop)
        m1.rebuild()
        objL = composite.read_stack(
                      stack = self.stack,
                      plyt  = self.plyt  )
        for ply in objL.plies:
            ply.matobj = m1
            ply.rebuild()
        objL.calc_lamination_parameters()
        self.params['a1'] = objL.xiA[1]
        self.params['a2'] = objL.xiA[2]
        self.params['a3'] = objL.xiA[3]
        self.params['a4'] = objL.xiA[4]
        self.params['b1'] = objL.xiB[1]
        self.params['b2'] = objL.xiB[2]
        self.params['b3'] = objL.xiB[3]
        self.params['b4'] = objL.xiB[4]
        self.params['d1'] = objL.xiD[1]
        self.params['d2'] = objL.xiD[2]
        self.params['d3'] = objL.xiD[3]
        self.params['d4'] = objL.xiD[4]

    def print_lam( self, log=None ):
        params = self.params.keys()
        params.sort()
        tmp = '['
        for i in range( len( self.stack ) ):
            theta = self.stack[i]
            if i == 0:
               tmp +=   '%3s' % str(theta)
            else:
               tmp += ', %3s' % str(theta)
        tmp += ' ]'
        tmp = '%50s' % tmp
        for key in params:
            value = '%1.6f' % self.params[ key ] 
            tmp  += '\t%10s' % value
        if log:
            log.write( tmp )
        else:
            print tmp
        return tmp
    
    def check_bloomfield( self ):
        import bloomfield
        bloomfield = reload( bloomfield )
        return bloomfield.check_params( self.params )

def print_lamination_parameters( lam_families ):
    num_plies = lam_families.keys()
    num_plies.sort()
    xiA = {1:[],2:[],3:[],4:[]}
    xiB = {1:[],2:[],3:[],4:[]}
    xiD = {1:[],2:[],3:[],4:[]}
    xiE = {1:[],2:[],3:[],4:[]}
    log = open(r'c:\users\pfh-castro\desktop\lam.txt', 'w')
    log.write( ('%50s' % 'plies') + '\tf_00\tf_90\tf_ang\tf_30\tf_45\tf_60\ta1\ta2\ta3\ta4\tb1\tb2\tb3\tb4\td1\td2\td3\td4\n' )
    for num_ply in num_plies:
        L = lam_families[ num_ply ]
        for dictL in L.values():
            pass
    log.close()
    for i in range(1,5):
        xiA[i].sort()
        xiB[i].sort()
        xiD[i].sort()
        xiE[i].sort()
    for i in range(1,5):
        print 'xi A %d\tMIN\t%f\tMAX\t%f' %( i, xiA[i][0], xiA[i][-1] )
    for i in range(1,5):
        print 'xi B %d\tMIN\t%f\tMAX\t%f' %( i, xiB[i][0], xiB[i][-1] )
    for i in range(1,5):
        print 'xi D %d\tMIN\t%f\tMAX\t%f' %( i, xiD[i][0], xiD[i][-1] )
    for i in range(1,5):
        print 'xi E %d\tMIN\t%f\tMAX\t%f' %( i, xiE[i][0], xiE[i][-1] )
    return True

